from operator import itemgetter
from typing import Union

from langchain_core.output_parsers import StrOutputParser
from langchain_core.output_parsers.openai_tools import JsonOutputToolsParser
from langchain_core.runnables import Runnable, RunnablePassthrough, RunnableLambda
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
# 加载 .env 到环境变量
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())


# 方法里面必须要描述才能被调用，比如下面的  """两个整数相乘"""
@tool
def multiply(first_int: int, second_int: int) -> int:
    """两个整数相乘"""
    return first_int * second_int


@tool
def add(first_int: int, second_int: int) -> int:
    "Add two integers."
    return first_int + second_int


@tool
def exponentiate(base: int, exponent: int) -> int:
    "Exponentiate the base to the exponent power."
    return base ** exponent


tools = [multiply, add, exponentiate]
tool_map = {tool.name: tool for tool in tools}


def call_tool(tool_invocation: dict) -> Union[str, Runnable]:
    """根据模型选择的 tool 动态创建 LCEL"""
    tool = tool_map[tool_invocation["type"]]
    return RunnablePassthrough.assign(
        output=itemgetter("args") | tool
    )


# .map() 使 function 逐一作用与一组输入
call_tool_list = RunnableLambda(call_tool).map()

# 模型
model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)


# 带有分支的 LCEL


def route(response):
    if len(response["functions"]) > 0:
        return response["functions"]
    else:
        return response["text"]


llm_with_tools = model.bind_tools(tools) | {
    "functions": JsonOutputToolsParser() | call_tool_list,
    "text": StrOutputParser()
} | RunnableLambda(route)

result = llm_with_tools.invoke("1024的平方是多少")
print(result)

result = llm_with_tools.invoke("你好")
print(result)
